/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: syyyy
 * Date: 2025-11-08
 * Time: 10:50
 */
import java.util.Scanner;

public class OptimizedOptimalBST {

    static class Node {
        int key;
        double prob;
        Node left, right;

        Node(int key, double prob) {
            this.key = key;
            this.prob = prob;
            this.left = this.right = null;
        }
    }

    /**
     * 优化的最优二叉搜索树算法（Knuth优化）
     * @param keys 排序后的关键字数组
     * @param probs 关键字搜索概率
     * @return 最优代价和根节点矩阵
     */
    public static Object[] optimizedOptimalBST(int[] keys, double[] probs) {
        int n = keys.length;
        double[][] cost = new double[n + 2][n + 2]; // 最优代价
        int[][] root = new int[n + 2][n + 2];       // 最优根节点
        double[][] sum = new double[n + 2][n + 2];  // 概率和
        int[][] range = new int[n + 2][n + 2];      // 记录分割点范围（优化用）

        // 初始化单个节点
        for (int i = 1; i <= n; i++) {
            cost[i][i] = probs[i - 1];
            root[i][i] = i;
            sum[i][i] = probs[i - 1];
            range[i][i] = i; // 单个节点的分割点只能是自己
        }

        // 按子树长度递增计算（从2到n）
        for (int len = 2; len <= n; len++) {
            for (int i = 1; i <= n - len + 1; i++) {
                int j = i + len - 1;
                sum[i][j] = sum[i][j - 1] + probs[j - 1]; // 计算概率和
                cost[i][j] = Double.MAX_VALUE;

                // 优化：分割点k的范围限制在[range[i][j-1], range[i+1][j]]
                int low = (j == i + 1) ? i : range[i][j - 1];
                int high = (j == i + 1) ? j : range[i + 1][j];

                // 只遍历[low, high]范围内的分割点
                for (int k = low; k <= high; k++) {
                    double currentCost = sum[i][j];
                    if (k > i) currentCost += cost[i][k - 1];
                    if (k < j) currentCost += cost[k + 1][j];

                    if (currentCost < cost[i][j]) {
                        cost[i][j] = currentCost;
                        root[i][j] = k;
                    }
                }
                range[i][j] = root[i][j]; // 记录当前区间的最优分割点，供后续范围限制用
            }
        }

        return new Object[]{cost[1][n], root};
    }

    // 构建二叉树（同之前的实现）
    private static Node buildTree(int[] keys, double[] probs, int[][] root, int i, int j) {
        if (i > j) return null;
        int k = root[i][j];
        Node node = new Node(keys[k - 1], probs[k - 1]);
        node.left = buildTree(keys, probs, root, i, k - 1);
        node.right = buildTree(keys, probs, root, k + 1, j);
        return node;
    }

    // 打印树结构（中序遍历，带缩进）
    private static void printTree(Node node, int depth) {
        if (node == null) return;
        printTree(node.right, depth + 1);
        for (int i = 0; i < depth; i++) System.out.print("    ");
        System.out.printf("Key: %d (概率: %.2f)\n", node.key, node.prob);
        printTree(node.left, depth + 1);
    }

    // 测试案例执行
    public static void runTestCase(int[] keys, double[] probs, int caseNum) {
        System.out.println("=== 测试案例 " + caseNum + " ===");
        System.out.print("关键字: ");
        for (int key : keys) System.out.print(key + " ");
        System.out.println();
        System.out.print("搜索概率: ");
        for (double p : probs) System.out.printf("%.2f ", p);
        System.out.println("\n---------------------");

        Object[] result = optimizedOptimalBST(keys, probs);
        double avgCost = (double) result[0];
        int[][] root = (int[][]) result[1];

        Node rootNode = buildTree(keys, probs, root, 1, keys.length);
        System.out.println("最优二叉搜索树结构:");
        printTree(rootNode, 0);
        System.out.printf("\n平均查询时间: %.4f\n", avgCost);
        System.out.println("=================================\n");
    }

    public static void main(String[] args) {
        // 测试案例1：3个关键字
        int[] keys1 = {10, 20, 30};
        double[] probs1 = {0.4, 0.3, 0.3};
        runTestCase(keys1, probs1, 1);

        // 测试案例2：5个关键字
        int[] keys2 = {1, 3, 5, 7, 9};
        double[] probs2 = {0.1, 0.2, 0.25, 0.3, 0.15};
        runTestCase(keys2, probs2, 2);

        // 测试案例3：6个关键字
        int[] keys3 = {2, 4, 6, 8, 10, 12};
        double[] probs3 = {0.05, 0.1, 0.15, 0.2, 0.25, 0.25};
        runTestCase(keys3, probs3, 3);

        // 用户自定义输入
        Scanner scanner = new Scanner(System.in);
        System.out.print("是否输入自定义数据？(y/n): ");
        if (scanner.next().equalsIgnoreCase("y")) {
            System.out.print("关键字数量: ");
            int n = scanner.nextInt();
            int[] keys = new int[n];
            double[] probs = new double[n];

            System.out.println("输入" + n + "个排序后的关键字:");
            for (int i = 0; i < n; i++) keys[i] = scanner.nextInt();
            System.out.println("输入对应的搜索概率（和为1）:");
            for (int i = 0; i < n; i++) probs[i] = scanner.nextDouble();

            runTestCase(keys, probs, 4);
        }
        scanner.close();
    }
}
